

import os
import cv2
from tqdm import tqdm

# 原始数据集路径
path = r"../datasets/proposed_Aug"
# CLAHE 处理后的数据保存路径
save_path = r"../datasets/ODIR-CLAHE"

# 确保保存路径存在
os.makedirs(save_path, exist_ok=True)

def clahe_filter(image):
    """对图像应用 CLAHE 滤波"""
    # 分离 BGR 三个通道
    b, g, r = cv2.split(image)

    # 创建 CLAHE 对象
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))

    # 应用 CLAHE 增强
    enhanced_b = clahe.apply(b)
    enhanced_g = clahe.apply(g)
    enhanced_r = clahe.apply(r)

    # 合并通道
    enhanced_image = cv2.merge([enhanced_b, enhanced_g, enhanced_r])
    return enhanced_image

if __name__ == "__main__":
    for root, dirs, files in os.walk(path):  # 遍历所有子文件夹
        for img_name in tqdm(files):
            img_path = os.path.join(root, img_name)  # 组合成完整路径
            input_image = cv2.imread(img_path)

            # 如果图片无法读取，则跳过
            if input_image is None:
                print(f"Warning: 无法读取 {img_path}")
                continue

            # 处理图像
            enhanced_image = clahe_filter(input_image)
            resized_image = cv2.resize(enhanced_image, (224, 224))

            # 计算对应的保存路径
            relative_path = os.path.relpath(root, path)  # 计算相对路径（子文件夹）
            target_folder = os.path.join(save_path, relative_path)
            os.makedirs(target_folder, exist_ok=True)  # 确保子文件夹存在
            
            # 保存增强后的图片
            save_img_path = os.path.join(target_folder, img_name)
            cv2.imwrite(save_img_path, resized_image)
